import sys
import os
from openai import OpenAI
import json
from typing import Optional

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

class DoubaoClient:
    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        self.client = OpenAI(
            api_key=api_key or "4b7ad9f8-24dc-4276-8d09-7fcc0cfab541",
            base_url=base_url or "https://ark.cn-beijing.volces.com/api/v3",
        )
    
    def build_message(self, text: str, image_urls: Optional[list] = None) -> list:
        """构建消息内容
        
        Args:
            text (str): 文本内容
            image_urls (list): 图片URL列表
            
        Returns:
            list: 消息内容
        """
        content = [{"type": "text", "text": text}]
        
        if image_urls:
            for url in image_urls:
                content.append({
                    "type": "image_url",
                    "image_url": {"url": url}
                })
                
        return [{
            "role": "user",
            "content": content
        }]
    
    def call_model(self, model: str, messages: list) -> dict:
        """调用大模型
        
        Args:
            model (str): 模型名称
            messages (list): 消息内容
            
        Returns:
            dict: 模型返回结果
        """
        response = self.client.chat.completions.create(
            model=model,
            messages=messages
        )
        return json.loads(response.choices[0].message.content)

    def save_result(self, result: dict, output_path: str) -> None:
        """保存结果到文件
        
        Args:
            result (dict): 结果数据
            output_path (str): 输出文件路径
        """
        with open(output_path, 'w', encoding='utf-8') as fout:
            fout.write(json.dumps(result, ensure_ascii=False, indent=2))

    def print_outline(self, result: dict) -> str:
        """生成试卷结构大纲
        
        Args:
            result (dict): 试卷结构JSON数据
            
        Returns:
            str: 试卷结构大纲字符串
        """
        if not result:
            return "无有效试卷数据"
            
        outline = []
        outline.append(f"试卷标题: {result.get('title', '未知')}")
        outline.append("\n试卷结构:")
        
        for section in result.get('sections', []):
            outline.append(f"\n{section.get('sectionNumber', '')}. {section.get('sectionTitle', '')} ({section.get('score', 0)}分)")
            
            for question in section.get('subQuestions', []):
                q_num = question.get('questionNumber', '')
                score = question.get('score', '')
                options = question.get('optionCount', '')
                
                line = f"  {q_num}. "
                if options:
                    line += f"选择题 ({options}个选项)"
                elif score:
                    line += f"分值: {score}分"
                    
                outline.append(line)
                
                # 处理子问题
                for sub_q in question.get('subSubQuestions', []):
                    outline.append(f"    {sub_q.get('questionNumber', '')}")
                    
        return "\n".join(outline)

# 测试代码
if __name__ == "__main__":
    client = DoubaoClient()
    
    text = '''这是一份考试的答题卡图片，请对图片进行OCR，并提取里面的试卷标题和题目结构。试卷标题包括大题号、小题号、分值。如果是选择题（客观题），除了给出题号外，还需要给出该小题的选项数量。请以json的形式输出试卷的题目结构，无需详细的题干内容。输出的json格式如下：{"title":"八年级教学质量监测数学答题卡","sections":[{"sectionNumber":"一","sectionTitle":"选择题","score":24,"subQuestions":[{"questionNumber":1,"optionCount":4}]},{"sectionNumber":"二","sectionTitle":"填空题","score":15,"subQuestions":[{"questionNumber":9}]},{"sectionNumber":"三","sectionTitle":"解答题","score":81,"subQuestions":[{"questionNumber":20,"score":5},{"questionNumber":21,"score":6,"subSubQuestions":[{"questionNumber":"(1)"},{"questionNumber":"(2)"}]}]}]}'''
    
    image_urls = [
        "https://yj-oss.yunxiao.com/v1/baidu-raw/template/9317293/2D24BA57-C195-11EF-A355-6C4B90985E68.png?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2025-01-24T06%3A32%3A35Z%2F3600%2Fhost%2F0bd3080e09ca49acfdf7c70897746c9b1219562dbb7d37df060f9772f7af978c",
        "https://yj-oss.yunxiao.com/v1/baidu-raw/template/9317293/2D24BA58-C195-11EF-A355-6C4B90985E68.png?authorization=bce-auth-v1%2Fa908715249bb41c998c7d924b2476b37%2F2025-01-24T06%3A32%3A35Z%2F3600%2Fhost%2Fcd84da3f4add6ab5fe672642f908211dbf720ad323874af199396320b1713727"
    ]
    
    messages = client.build_message(text, image_urls)
    try:
        result = client.call_model("ep-20250124123709-scsq6", messages)
        client.save_result(result, "./output/json_out.json")
        outline = client.print_outline(result)
        print("试卷结构：")
        print(outline)
    except Exception as e:
        print(f"Error: {str(e)}")
